import traceback

import websocket
import threading
import time
import json
import sys
import ChatMessage_pb2
USE_PROTOBUF = False


def on_message(ws, message):
    if isinstance(message, bytes):
        chat_msg = ChatMessage_pb2.ChatMessage()
        try:
            chat_msg.ParseFromString(message)
            print(f"📩 收到来自 {chat_msg.sender} 的消息: {chat_msg.content}")
        except Exception as e:
            print(f"❌ Protobuf 解码失败: {e}")
    else:
        print(f"⚠️ 收到非二进制消息: {message}")

def on_error(ws, error):
    print(f"连接错误: {error}")


def on_close(ws, close_status_code, close_msg):
    print("连接关闭")


def on_open(ws, user_id, receiver_id):
    print(f"✅ {user_id} 连接成功")

    def run():
        print(f"请输入消息，按回车发送给 {receiver_id}，输入 exit 退出")
        while True:
            try:
                msg = sys.stdin.readline()
                if msg == "":
                    continue
                msg = msg.strip()
                if msg.lower() == "exit":
                    print("👋 正在关闭连接...")
                    ws.close()
                    break
                if not msg:
                    continue

                timestamp = int(time.time())

                if USE_PROTOBUF:
                    # 构造 Protobuf 消息对象
                    message = ChatMessage_pb2.ChatMessage(
                                            sender=user_id,
                                            receiver=receiver_id,
                                            content=msg,
                                            timestamp=timestamp
                                            )
                    ws.send(message.SerializeToString(), opcode=websocket.ABNF.OPCODE_BINARY)
                else:
                    # JSON 发送
                    message = {
                        "sender": user_id,
                        "receiver": receiver_id,
                        "content": msg,
                        "timestamp": timestamp,
                    }
                    ws.send(json.dumps(message))  # 文本方式发送

                print(f"你发送了: {msg}")
            except Exception as e:
                print(f"❌ 发送消息失败: {e}")
                traceback.print_exc()
                break

    # 使用守护线程，程序退出时自动结束
    threading.Thread(target=run, daemon=True).start()


if __name__ == "__main__":
    user_id = "userA"
    receiver_id = "userB"
    jwt_token = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJ1c2VyQSIsImlhdCI6MTc1MjA1NDkzMywiZXhwIjoxNzUyMDU4NTMzfQ.26JV6K9SkgJlU1G6Jp_jJHTLASb8CkoWcFvB7c4OvvE"

    ws_url = "ws://localhost:9090/ws"
    headers = [f"Cookie: jwt={jwt_token}"]

    ws = websocket.WebSocketApp(ws_url,
                                header=headers,
                                on_message=on_message,
                                on_error=on_error,
                                on_close=on_close,
                                on_open=lambda ws: on_open(ws, user_id, receiver_id))

    ws.run_forever()
